import json
import os
import pickle
import zipfile

import joblib
import torch
from loguru import logger
from typing import Dict, Any, Optional, Tuple, Union
import io


class ModelCreatorConfig:
    def __init__(self,
                 suffix: Optional[str] = None,
                 model_path: Optional[str] = None,
                 model_name: str = "model",
                 pth_path: Optional[str] = None,
                 static_pth_path: Optional[str] = None,
                 pkl_path: Optional[str] = None):
        self.features_columns = []
        self.pth_path = pth_path
        self.static_pth_path = static_pth_path
        self.pkl_path = pkl_path
        self.model_name = model_name
        self.model_path = model_path
        suffix = suffix or "mod"
        model_name = model_name or "model"
        model_path = model_path or os.getcwd()
        self.model_save_path = os.path.join(model_path, f"{model_name}.{suffix}")
        # 确保目录存在
        for path in [self.pth_path, self.static_pth_path, self.pkl_path]:
            if path and not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
                logger.warning(f"路径 {os.path.dirname(path)} 不存在，已自动创建")


class ModelCreator:
    """
        模型创建器
        用于创建模型并保存模型
        """

    def __init__(self, model_config: ModelCreatorConfig):
        """
        初始化模型创建器
        :param model_config: 模型创建器配置
        """
        self.model_config = model_config

    def create(self) -> str:
        logger.info(f"开始创建模型 {self.model_config.model_name}")

        # 读取模型数据
        with open(self.model_config.pth_path, 'rb') as f:
            pth_content = f.read()

        with open(self.model_config.static_pth_path, 'rb') as f:
            static_pth_content = f.read()

        with open(self.model_config.pkl_path, 'rb') as f:
            pkl_content = f.read()

        # 确保保存目录存在
        os.makedirs(os.path.dirname(self.model_config.model_save_path), exist_ok=True)

        # 使用pickle保存
        with zipfile.ZipFile(self.model_config.model_save_path, 'w') as z:
            z.writestr('model.pth', pth_content)
            z.writestr('static_model.pth', static_pth_content)
            z.writestr('scaler.pkl', pkl_content)
            z.writestr('features.json', json.dumps(self.model_config.features_columns).encode('utf-8'))

        logger.info(f"模型 {self.model_config.model_name} 保存成功")
        return self.model_config.model_save_path

    def remove_temp(self):
        try:
            # 删除PTH和PKL文件
            os.remove(self.model_config.pth_path)
            os.remove(self.model_config.static_pth_path)
            os.remove(self.model_config.pkl_path)
            # 删除文件夹
            if os.path.exists(os.path.dirname(self.model_config.static_pth_path)):
                os.rmdir(os.path.dirname(self.model_config.static_pth_path))
                logger.debug(
                    f"模型{self.model_config.model_name}静态权重文件{self.model_config.static_pth_path}删除成功")
            if os.path.exists(os.path.dirname(self.model_config.pth_path)):
                os.rmdir(os.path.dirname(self.model_config.pth_path))
                logger.debug(f"模型{self.model_config.model_name}权重文件{self.model_config.pth_path}删除成功")
            if os.path.exists(os.path.dirname(self.model_config.pkl_path)):
                os.rmdir(os.path.dirname(self.model_config.pkl_path))
                logger.debug(f"模型{self.model_config.model_name}参数文件{self.model_config.pkl_path}删除成功")
            logger.info(f"模型{self.model_config.model_name}临时缓存文件删除成功")
            return True
        except OSError as e:
            logger.error(f"模型{self.model_config.model_name}临时缓存文件删除失败，原因{e}")
            return False


class ModelLoader:
    def __init__(self, model_path: str, device: Optional[str] = None):
        self.model_path = model_path
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.features_columns = []
        self.model = None
        self.scaler = None
        self.static_model = None

    def load(self) -> Tuple[list, Any, Any]:
        # 读取pickle数据
        try:
            with zipfile.ZipFile(self.model_path, 'r') as z:
                with z.open('model.pth') as f:
                    self.model = torch.load(f, map_location=self.device, weights_only=False)
                with z.open('static_model.pth') as f:
                    self.static_model = torch.load(f, map_location=self.device, weights_only=True)
                with z.open('scaler.pkl') as f:
                    self.scaler = joblib.load(f)
                with z.open('features.json') as f:
                    self.features_columns = json.loads(f.read().decode('utf-8'))

                return self.features_columns, self.model, self.scaler
        except Exception as e:
            logger.error(f"模型 {self.model_path} 加载失败: {e}")
            raise IOError(f"模型 {self.model_path} 加载失败: {e}")
